package com.wfh.frp.handler;

import com.wfh.frp.model.ProxyMsg;

import io.netty.channel.Channel;

import io.netty.channel.ChannelHandlerContext;

import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.timeout.IdleState;
import io.netty.handler.timeout.IdleStateEvent;

import lombok.extern.slf4j.Slf4j;

import java.util.Map;

import java.util.concurrent.ConcurrentHashMap;

/**
 * @Title: ServerHandler
 * @Author wangfenghuan
 * @Package com.wfh.frp.handler
 * @Date 2025/10/8 11:05
 * @description: 主服务处理器
 */
@Slf4j
public class ServerHandler extends SimpleChannelInboundHandler<ProxyMsg> {

    /**
     * 存储访客通道
     */
    private static final Map<String, Channel> visitorChannel = new ConcurrentHashMap<>();

    /**
     * 存储内网服务通道
     */
    private static final Map<String, Channel> serviceChannel = new ConcurrentHashMap<>();

    /**
     * 心跳间隔（秒）
     */
    private static final int HEARTBEAT_INTERVAL = 30;

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
        if (evt instanceof IdleStateEvent) {
            IdleStateEvent event = (IdleStateEvent) evt;
            if (event.state() == IdleState.WRITER_IDLE) {
                // 发送心跳消息
                ProxyMsg heartbeat = new ProxyMsg();
                heartbeat.setType(ProxyMsg.HEARTBEAT);
                ctx.writeAndFlush(heartbeat);
                log.debug("Sent heartbeat to client: {}", ctx.channel().id());
            }
        } else {
            super.userEventTriggered(ctx, evt);
        }
    }

    @Override
    protected void messageReceived(ChannelHandlerContext channelHandlerContext, ProxyMsg msg) throws Exception {
        byte type = msg.getType();
        String channelId = msg.getChannelId();

        if (type == ProxyMsg.SERVICE_CONNECT) {
            if (serviceChannel.containsKey(channelId)) {
                channelHandlerContext.close();
                log.warn("Service channel {} is already registered, closing new connection", channelId);
            } else {
                serviceChannel.put(channelId, channelHandlerContext.channel());
                log.info("Service channel registered: {}", channelId);
            }
        } else if (type == ProxyMsg.VISITOR_CONNECT) {
            if (visitorChannel.containsKey(channelId)) {
                channelHandlerContext.close();
                log.warn("Visitor channel {} is already registered, closing new connection", channelId);
            } else {
                visitorChannel.put(channelId, channelHandlerContext.channel());
                log.info("Visitor channel registered: {}", channelId);
            }
        } else if (type == ProxyMsg.TRANSFER) {
            if (msg.getData() != null) {
                Channel currentChannel = channelHandlerContext.channel();

                // 如果是访客发的 → 转发给内网
                if (visitorChannel.containsValue(currentChannel)) {
                    Channel serviceCh = serviceChannel.get(channelId);
                    if (serviceCh == null || !serviceCh.isActive()) {
                        currentChannel.close();
                        log.warn("Service channel for {} is not active, closing visitor channel", channelId);
                    } else {
                        serviceCh.writeAndFlush(msg);
                    }
                }
                // 如果是内网发的 → 转发给访客
                else if (serviceChannel.containsValue(currentChannel)) {
                    Channel visitorCh = visitorChannel.get(channelId);
                    if (visitorCh == null || !visitorCh.isActive()) {
                        currentChannel.close();
                        log.warn("Visitor channel for {} is not active, closing service channel", channelId);
                    } else {
                        visitorCh.writeAndFlush(msg);
                    }
                }
            }
        } else if (type == ProxyMsg.DISCONNECT) {
            Channel v = visitorChannel.remove(channelId);
            Channel s = serviceChannel.remove(channelId);
            if (v != null) {
                v.close();
            }
            if (s != null) {
                s.close();
            }
        }
    }

    /**
     * 通道关闭
     * @param ctx
     */
    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        Channel channel = ctx.channel();
        String channelId = null;

        // Find and remove from visitor channels
        for (Map.Entry<String, Channel> entry : visitorChannel.entrySet()) {
            if (entry.getValue() == channel) {
                channelId = entry.getKey();
                visitorChannel.remove(entry.getKey());
                log.info("Visitor channel disconnected: {}", channelId);
                break;
            }
        }

        // Find and remove from service channels
        for (Map.Entry<String, Channel> entry : serviceChannel.entrySet()) {
            if (entry.getValue() == channel) {
                channelId = entry.getKey();
                serviceChannel.remove(entry.getKey());
                log.info("Service channel disconnected: {}", channelId);
                break;
            }
        }

        // If we found a channel ID, send disconnect message to other side
        if (channelId != null) {
            ProxyMsg disconnectMsg = new ProxyMsg();
            disconnectMsg.setType(ProxyMsg.DISCONNECT);
            disconnectMsg.setChannelId(channelId);

            // Forward disconnect message to the other side if they're still connected
            Channel otherChannel = visitorChannel.get(channelId);
            if (otherChannel == null) {
                otherChannel = serviceChannel.get(channelId);
            }

            if (otherChannel != null && otherChannel.isActive()) {
                otherChannel.writeAndFlush(disconnectMsg);
            }
        }
    }
}